
#include "book.h"

/*
void save_idf(const std::string& filename, const std::map<size_t,double>& idf) {

    Sink fout(filename);

    fout << idf.size();

    for (const auto& z : idf) {
        fout << z.first;
        fout << z.second;
    }
}

*/

void read_terms_for_idf(const std::string& filename, std::map<size_t,double>& out) {

    Source fin(filename);

    std::cerr << "Reading terms for idf... " << filename << std::endl;

    while (1) {

        std::map<size_t,size_t> tmp;

        if (!read_terms(tmp, fin)) {
            break;
        }

        for (const auto& z : tmp) {
            out[z.first] += z.second;
        }
    }
}

void process_one_book(Book& book, std::map<size_t,double>& idf, const std::string& outfilename) {

    std::map<size_t,double> tfidf;

    double sum = 0;
    for (const auto& z : book.w) {

        double v;

        // HACK
        if (z.second == idf[z.first]) {
            v = 0;

        } else {

            v = (double)z.second / ((double)idf[z.first] - z.second + 1);
        }

        tfidf[z.first] = v;
        sum += (v * v);
    }

    sum = ::sqrt(sum);

    std::map< double,std::vector<size_t> > resort;

    for (auto& z : tfidf) {
        z.second /= sum;
        resort[z.second].push_back(z.first);
    }

    //

    tfidf.clear();

    size_t nj = 0;
    auto z = resort.rbegin();

    static const size_t N = 300;

    while (z != resort.rend()) {

        for (const auto& y : z->second) {
            tfidf[y] = z->first;
            ++nj;

            if (nj >= N)
                break;
        }

        if (nj >= N)
            break;

        ++z;
    }

    //

    write_digest(book, tfidf, outfilename);

}

void read_terms_for_tf(const std::string& filename, std::map<size_t,double>& idf, const std::string& outfilename) {

    {
        Sink fout(outfilename);
    }

    Source fin(filename);

    std::cerr << "Reading terms for tf... " << filename << std::endl;

    while (1) {

        Book tmp;

        if (!read_book(tmp, fin)) {
            break;
        }

        process_one_book(tmp, idf, outfilename);
    }
}

int main(int argc, char** argv) {
    
    try {

        std::map<size_t,double> idf;

        for (int i = 1; i < argc; ++i) {

            std::string books(argv[i]);

            read_terms_for_idf(books + ".terms", idf);
        }

        for (int i = 1; i < argc; ++i) {

            std::string books(argv[i]);
            read_terms_for_tf(books, idf, books + ".tfidf");
        }

    } catch (std::exception& e) {
        std::cerr << "ERROR: " << e.what() << std::endl;
        return 1;

    } catch (...) {
        std::cerr << "UNKNOWN ERROR" << std::endl;
        return 1;
    }

    return 0;
}


